package com.atguigu.state;

import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

/**
 * @author gmd
 * @desc 算子状态-列表状态。在map算子中计算数据的个数
 * @since 2024-12-01 11:31:52
 */
public class OperatorListState {

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(2);

        env.socketTextStream("127.0.0.1", 7777)
                .map(new MyCountMapFunction())
                .print();

        env.execute();
    }


    // 1.实现 CheckpointedFunction 接口
    public static class MyCountMapFunction implements MapFunction<String, Long>, CheckpointedFunction {

        private Long count = 0L;
        private ListState<Long> state;


        @Override
        public Long map(String value) throws Exception {
            return ++count;
        }

        /*
         * 2.本地变量持久化：将 本地变量 拷贝到 算子状态中。
         * 该方法在开启checkpoint时才会调用，用于保存当前的状态。
         *
         * @param context FunctionSnapshotContext对象，提供快照相关的上下文信息
         */
        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            // 在快照状态之前，打印日志以标记该操作
            System.out.println("状态保存（snapshotState）...");
            // 2.1 清空算子状态，以确保状态的一致性
            state.clear();
            // 2.2 将 本地变量 添加到 算子状态 中，实现状态的持久化
            state.add(count);
        }

        /*
         * 3.初始化本地变量：程序启动和恢复时， 从状态中 把数据添加到 本地变量，每个子任务调用一次。
         * 该方法在程序启动和恢复时被调用，用于从状态中恢复本地变量的值。
         *
         * @param context FunctionInitializationContext对象，提供初始化相关的上下文信息
         */
        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            // 标记初始化操作
            System.out.println("状态初始化（initializeState）...");

            // 3.1 从上下文初始化算子状态
            // 获取算子状态存储，并从中获取列表状态
            state = context.getOperatorStateStore()
                    .getListState(new ListStateDescriptor<Long>("state", Types.LONG));
            // 联合列表状态
            // .getUnionListState(new ListStateDescriptor<Long>("union-state", Types.LONG));

            // 3.2 从算子状态中把数据拷贝到本地变量
            // 如果程序是从故障中恢复的，则需要从状态中恢复本地变量的值
            if (context.isRestored()) {
                // 遍历状态中的值，并累加到本地变量中
                for (Long c : state.get()) {
                    count += c;
                }
            }
        }
    }

    /*
     * 算子状态中， list 与 unionlist 的区别：
     * 并行度改变时，怎么重新分配状态。
     * 两种状态的主要区别在于并行度改变时的状态分配策略：
     *
     * 1. List 状态：轮询均分给新的并行子任务
     *   - 当并行度增加时，原有的状态会被均分给新的子任务
     *   - 当并行度减少时，多个子任务的状态会被合并
     *
     * 2. UnionList 状态：原先的多个子任务的状态，合并成一份完整的
     *   - 当并行度改变时，原有的状态会被合并成一份完整的列表
     *   - 新的并行子任务会接收到完整的列表（每人一份完整的）
     */

}
